import os
import glob
from typing import List, Callable

from PIL import Image
from torchvision import transforms
from torch.utils.data import DataLoader, Dataset

THIS_DIR = os.path.dirname(os.path.abspath(__file__))

LABELS = [
    "airplane", "automobile", "bird", "cat", "deer",
    "dog", "frog", "horse", "ship", "truck",
]

LABEL_DICT = dict(zip(LABELS, range(len(LABELS))))


def default_loader(im_filepath: str):
    return Image.open(im_filepath).convert("RGB")


# train_transform = transforms.Compose(
#     [
#         transforms.RandomResizedCrop((28, 28)),  # 随机裁剪之后，将图像resize到28*28 原来是32*32
#         transforms.RandomHorizontalFlip(),  # 水平翻转，概率默认0.5
#         transforms.RandomVerticalFlip(),  # 垂直翻转，概率默认0.5
#         transforms.RandomRotation(90),  # -90~90旋转
#         transforms.RandomGrayscale(0.1),  # 灰度,以0.1的概率置灰
#         transforms.ColorJitter(0.3, 0.3, 0.3, 0.3),
#         transforms.ToTensor(),
#     ]
# )

# 稍微写简单一些
train_transform = transforms.Compose([
    transforms.RandomCrop(28),  # 随机裁剪
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
])

test_transform = transforms.Compose([
    transforms.Resize((28, 28)),
    transforms.ToTensor(),
])


class MyDataset(Dataset):

    def __init__(
            self,
            im_list: List[str],
            transform: Callable = None,
            loader: Callable = default_loader
    ):
        """
        初始化函数。
        这里只是示例，我们理解后可以传入任何形式的参数，只要能完成文件/数据的加载
        :param im_list: 图片文件列表, 是每一张图片的全路径
        :param transform: transform函数
        :param loader: 图片加载器，如何将一张一张图片加载为矩阵，需要自己实现
        """
        super().__init__()
        self.imgs: List[List[str, int]] = []  # 每个元素为：[图片路径，图片所属类别数字]

        for im_item in im_list:
            # D:\CodingFiles\PycharmProjects\pytorchfirst\06.cifar10_cls\datasets\TRAIN\airplane\aeroplane_s_000004.png
            _items = im_item.split(os.path.sep)
            im_label_name = _items[-2]
            self.imgs.append([im_item, LABEL_DICT[im_label_name]])

        self.transform = transform
        self.loader = loader

    def __getitem__(self, index):
        im_path, im_label = self.imgs[index]

        im_data = self.loader(im_path)

        # 如果有数据增强，在这里调用数据增强。。所以transform的输入和数据息息相关
        # 数据增强，在训练和测试的时候是不一样的。主要是训练的时候，需要加入数据增强
        if self.transform is not None:
            im_data = self.transform(im_data)

        return im_data, im_label

    def __len__(self):
        return len(self.imgs)


# if __name__ == '__main__':
im_train_list = glob.glob(rf"{THIS_DIR}\datasets\TRAIN\*\*.png")
im_test_list = glob.glob(rf"{THIS_DIR}\datasets\TEST\*\*.png")

# print("im_train_list: ", im_train_list)
# print("im_test_list: ", im_test_list)

train_dataset = MyDataset(im_train_list, transform=train_transform, loader=default_loader)
# test_dataset = MyDataset(im_test_list, transform=transforms.ToTensor(), loader=default_loader)
test_dataset = MyDataset(im_test_list, transform=test_transform, loader=default_loader)

train_data_loader = DataLoader(
    dataset=train_dataset,
    batch_size=128,
    shuffle=True,
    num_workers=0,  # windows下并行有坑，改为1个cpu
)
test_data_loader = DataLoader(
    dataset=test_dataset,
    batch_size=128,
    shuffle=False,
    num_workers=0,
)

print("num of train: ", len(train_dataset))
print("num of test: ", len(test_dataset))
